import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain.memory import ConversationSummaryBufferMemory
from langchain.chains import LLMChain
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

# --- 环境准备 ---
print("正在从 env.example 加载环境变量...")
load_dotenv(dotenv_path="env.example")
api_key = os.getenv("OPENAI_API_KEY")
model_name = os.getenv("OPENAI_MODEL", "qwen-plus-latest")

if not api_key:
    print("错误：请先在 env.example 文件中设置你的 OPENAI_API_KEY。")
else:
    print("环境变量加载成功！")
    # --- 核心代码 (使用经典组件以突出摘要功能) ---

    # 1. 初始化模型
    llm = ChatOpenAI(model=model_name)
    print(f"模型已初始化 (使用 {model_name})。")

    # 2. 定义摘要记忆模块
    # 这是本节课的明星！
    # - llm: 我们需要传递一个llm实例，因为它需要调用AI来生成摘要。
    # - max_token_limit: 一个关键参数。当历史记录的Token数超过这个限制时，
    #   它就会触发一次摘要，将最旧的对话压缩成摘要，并保留最新的对话在缓冲区。
    #   我们把它设得很小（比如 200 Tokens），以便在示例中能轻易触发它。
    memory = ConversationSummaryBufferMemory(
        llm=llm,
        max_token_limit=200,
        memory_key="chat_history",
        return_messages=True
    )
    print("摘要记忆模块已初始化 (max_token_limit=200)。")

    # 3. 创建提示模板 (与之前相同)
    prompt = ChatPromptTemplate(
        input_variables=["input"],
        messages=[
            MessagesPlaceholder(variable_name="chat_history"),
            ("user", "{input}")
        ]
    )

    # 4. 创建 LLMChain (为了清晰演示摘要记忆，我们暂时使用经典链)
    chain = LLMChain(
        llm=llm,
        prompt=prompt,
        memory=memory,
        verbose=True  # verbose=True 在这里至关重要！
    )
    print("LLMChain 已创建。")
    print("\n--- 注意：verbose=True 将打印出发送给模型的最终提示。---")
    print("--- 请仔细观察对话变长时，提示内容如何从完整历史变为[摘要+最新历史] ---")

    # 5. 开始一场足以触发摘要的长对话
    # 第一轮
    print("\n[你]: 你好，我叫派大星，我的朋友是海绵宝宝，我们住在比奇堡。")
    chain.invoke({"input": "你好，我叫派大星，我的朋友是海绵宝宝，我们住在比奇堡。"})

    # 第二轮
    print("\n[你]: 能给我介绍一下 LangChain 框架的核心思想吗？请详细一点说。")
    chain.invoke({"input": "能给我介绍一下 LangChain 框架的核心思想吗？请详细一点说。"})

    # 第三轮 - 这一轮很可能会触发摘要！
    print("\n[你]: 那它的 Prompt Templates 和 Output Parsers 具体有什么用呢？")
    chain.invoke({"input": "那它的 Prompt Templates 和 Output Parsers 具体有什么用呢？"})
    
    # 第四轮 - 验证记忆
    print("\n[你]: 你还记得我叫什么，我的朋友是谁吗？")
    response = chain.invoke({"input": "你还记得我叫什么，我的朋友是谁吗？"})

    print(f"\n[AI 最终回答]: {response['text']}")
    print("\n--- 如果 AI 能回答正确，说明摘要记忆成功保留了早期核心信息！ ---")
    
    print("\n--- 查看 Memory 内部状态 ---")
    print(memory.load_memory_variables({}))
    print("--------------------") 